import os
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset,random_split
import pandas as pd
import PIL.Image as Image
from torchvision import transforms as transform
import re

classification_names = {
    0: '上身衣服',  # 数字 0 对应“上身衣服”
    1: '鞋',       # 数字 1 对应“鞋”
    2: '包',       # 数字 2 对应“包”
    3: '下身衣服',  # 数字 3 对应“下身衣服”
    4: '手表'      # 数字 4 对应“手表”
}

#定义函数，对图片排序
def sort_alphanum(img_names):
    #数字转Int
    convert = lambda text: int(text) if text.isdigit() else text.lower()
    aplhanum_key = lambda img_names: [convert(c) for c in re.split('([0-9]+)', img_names)]
    return sorted(img_names, key=aplhanum_key)

class ImageDataset(Dataset):
    def __init__(self, image_dir, transform=None):
        super().__init__()
        self.image_dir = image_dir
        self.transform = transform
        self.all_images = sort_alphanum(os.listdir(image_dir))
        self.labels = pd.read_csv('../dataset/fashion-labels.csv')  
        self.labels_dict = dict(zip(self.labels['id'],self.labels['target']))

    def __len__(self):
        return len(self.all_images)

    def __getitem__(self, idx):
        img_path = os.path.join(self.image_dir, self.all_images[idx])
        image = Image.open(img_path).convert("RGB")
        if self.transform:
            image = self.transform(image)
        label = self.labels_dict[idx]
        
        #返回（噪声图片，原始图片）
        return image,label
    
transform = transform.Compose([
    transform.Resize((64,64)),
    transform.ToTensor()])
full_dataset = ImageDataset(image_dir='../dataset',transform=transform)
print(len(full_dataset.labels))

#划分数据集
train_dataset, test_dataset = torch.utils.data.random_split(full_dataset, [0.75, 0.25])

train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=32, shuffle=True,drop_last=True)

test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=32, shuffle=False)

class Classifier(nn.Module):
    def __init__(self, in_channels=1, num_classes=5):
        super(Classifier, self).__init__()
        self.model = nn.Sequential(
        # 定义第一个卷积层，输入通道数为3，输出通道数为8，卷积核大小为3x3，步幅为1，填充为1
        nn.Conv2d(in_channels=3, out_channels=8, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
        # 使用填充为1的卷积操作，输出特征图的尺寸与输入相同（Same convolutions）
        nn.ReLU(),  # 定义ReLU激活函数
        # 定义最大池化层，池化核大小为2x2，步幅为2
        nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 2)),

        # 定义第二个卷积层，输入通道数为8，输出通道数为16，卷积核大小为3x3，步幅为1，填充为1
        nn.Conv2d(in_channels=8, out_channels=16, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)),
        # 上一层的输出通道数为8，因此这一层的输入通道数为8
        nn.ReLU(),  # 定义ReLU激活函数
        # 定义最大池化层，池化核大小为2x2，步幅为2
        nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 2)),
        # 定义全连接层，输入大小为16*16*16，输出大小为num_classes（分类数）
        nn.Flatten(),
        nn.Linear(16*16*16, num_classes)
        # 输入大小为16是因为conv2的输出通道数为16，16*16是因为经过两次池化后，特征图的尺寸为16x16
        )


    def forward(self, x):
        return self.model(x)
